查看原文
其他

TensorFlow Hub:可重用机器学习模型库

Google TensorFlow 2021-07-27

我们知道在机器学习中,模型训练需要大量时间与计算资源。为了提升资源利用与效率 我们创建了 TensorFlow Hub 项目,旨在促进机器学习中可复用模型的再利用 —— 发布、探索与使用的代码库。


这个模型库中包含了独立的 TF 计算图、权重及数据,开发者可以将 TensorFlow 的预训练模型片段重用在新的任务上。


用户可以在不同任务间进行重用,如:
  • 使用小型数据集进行模型训练
  • 提升泛化能力
  • 提升训练速度


我们在去年开源了 TensorFlow Hub,经过一年多的快速发展,我们为大家带来了项目的最新进展:



何时应该使用 TF Hub?

重新训练一个模型需要海量数据,如果你还在纠结于怎么收集到足够多的数据,不如利用迁移学习:使用 TF Hub 轻松重用一些已经过大量数据集训练过的模型。


通过重用模型可以节省海量的计算时间与计算资源,利用 TF Hub,无需花费太多时间进行学习与训练,即可快速将模型部署至神经网络中。只需一行代码即可调用图像、文本、视频等模型。


你可能还会面临一些其他问题,如:代码库耦合程度高,或者实验进展愈加缓慢……就像充分利用代码存储库可以让软件工程师的开发效率更高一样,通过定义不依赖于代码的 Artifact,TF Hub 让系统维护更轻松。


TF Hub 的核心理念是预训练构建模块通常,我们首先选取适当的算法和数据集来训练大型模型。我们可能只对模型中的一部分感兴趣 (可重用部分),如 Bottleneck Layer 或 可重用分布式表示 (Reusable Distributed Representation)。


使用 SavedModel 进行封装,SavedModel 保存了模型的完整信息(如参数的权重与计算图),封装后的 Artifact 将不再依赖于原始代码,这将更有利于模型的分享与部署。您可以通过网络或云进行共享,或在新的任务中直接部署。


然后,你可以将此模块作为新任务模型的一部分,模块由  TensorFlow Primitives 定义,你可对权重进行微调,以适应 新任务中 要解决问题。



在 TF2.0 中 TF Hub 有哪些新功能?

在过去的几个月里,我们致力于增强 TF Hub 的易用性。我们逐步将“序列化网络模型的一部分并将其重新加载使用”的理念融入 TensorFlow 核心之中,增强了 SavedModel 的功能,从而便于分享。通过 Eager Execution,选择网络中的需导出的部分更为简单易用。


让我们来看几个例子:import tensorflow_hub as hub
import numpy as np
# TF-2.0 and eager
text_features = hub.load(
"https://hub.tensorflow.google.cn/google/tf2-preview/nnlm-zh-dim128/1")
example_embeddings = text_features(np.array(["你好", "世界"]))


# SavedModel can share TF primitives such as tf.function and tf.Variables.
text_features.__call__
Text_features.embeddings


在 TensorFlow 2.0 中,我们使用 hub.load 加载模块。如果模型已经在我们的文件系统中,我们还能使用 tf.savedmodel.load 进行加载。通过即时运行功能,加载后即可调用。


借助新功能,您现在可以分享任何使用 TensorFlow Primitives 编写的对象。如在本例中 text_features 的两个成员:_call_(tf.functions) 与 embeddings ( tf.variables )。TF Hub 包含 TensorFlow 2.0 中支持多态函数的 SavedModels。  我们也在本次更新中提供了比上一版本更为直观的接口:import tensorflow_hub as hub


# Polymorphic functions
embed = hub.load(
"https://hub.tensorflow.google.cn/google/tf2-preview/mobilenet_v2/feature_vector/1")


embed(images, training=False)
embed(images, training=True)
embed(images, training=True, batch_norm_momentum=0.98)


图像表示模块在加载后进行推断,或在训练模式中采用了 BatchNorm。此时,我们通过 tf.function 控制部分参数(参数已包含在 TF 计算图中,无需再四处调用,并且这些 API 看起来都很规整):

import tensorflow_hub as hub
import tensorflow as tf
from tensorflow import keras
import numpy as np



# Integration with Keras
hub_layer = hub.KerasLayer(
"https://hub.tensorflow.google.cn/google/tf2-preview/nnlm-zh-dim128/1",
output_shape=[128], input_shape=[], dtype=tf.string)


model = keras.Sequential()
model.add(hub_layer)
model.add(keras.layers.Dense(16, activation='relu'))
model.add(keras.layers.Dense(1, activation='sigmoid'))


model.summary()
print(model.predict(np.array(["你好", "世界", "1, 2, 3"])))


我们还加入了一个新函数:hub.KerasLayer,在 Keras 模型集成  Hub 模块变得更为简单。


在本例中,我们构建了一个语句分类模型,模型共有三层。顶层是 Keras 层 (Neural-Net Language Model, NNLM),把句子作为输入数据,输出稠密向量编码,与全连接层和分类层。

  • Neural-Net Language Model
    http://www.jmlr.org/papers/volume3/bengio03a/bengio03a.pdf


Keras 层中包含文本预处理功能,我们可以直接输入语料,无需再定义预处理的逻辑。如果想尝试其他文本模型也很简单,修改 URL 即可切换模型。


在 TensorFlow Hub 0.3 版本以上有两个新函数 —— hub.load 和 hub.KerasLayer。新功能在 TensorFlow 2 中亦可使用,支持 Eager 与 Graph Mode。


我们已发布几个采用此模块的实例,您可预览此功接下来,我们将着手向后兼容已有模块。



常用模块介绍

目前,Google AI 和 DeepMind 团队已在 TF Hub 上分享了他们使用的一些模型。这些项目在去年便已发布,部分模块在近期更新。我们将介绍部分常用模块:


通用语句编码器


通用语句编码器是社区中最受欢迎的模块,将短句编码为多维向量,可用于多种自然语言任务。


最近,我们增加了它的多语言版本。因此,无论是哪种语言,只要是意义相近似的句子,最终都会出现在相邻的点上。在新版本中,你可以使用英语数据集对分类器进行训练,然后使用其他语言继续训练该模型。


图像增强模块

我们还加入了图像增强模块。增强图像的策略源于一些强化学习的训练任务(如 ImageNet),这些策略均已经验证,可迁移至新任务中应用。您也可以直接使用此模块或图像表示模块,然后通过字符串把这些模块串联在一起。


在本例中,图像增强模块可通过数据增强而减少需求的数据量,并能通过迁移学习减少图像特征向量。


新版本集成了很多有趣的功能,如:
  • 应用于文本任务的 BERT 模块

  • 对象检测模块

  • 控制图像生成的 BigGAN 模块

  • 视频动作识别的 I3D Kinetics

……还有更多等你来探索 (部分模型针对低资源而设计)。


我们期望让 TF Hub 与 TensorFlow 生态系统中其他部分的集成度更高,如命令行实用程序,将 TF Hub 模块转换为 tfjs 模块;让 TF Hub 模块可与 AutoML 中的代码库 AdaNet 一起使用;在 TF Transform 中使用 TF Hub 模块等。


如果你想尝试使用 TF Hub,请前往 hub.tensorflow.google.cn 进行了解,通过搜索对应模块及学习示例了解如何使用。



如果您想详细了解 本文提及 的相关内容,请参阅以下文档。这些文档深入探讨了这篇文章中提及的许多主题:



    您可能也对以下帖子感兴趣

    文章有问题?点此查看未经处理的缓存